import torch
from torchvision import transforms
from PIL import Image
import numpy as np
import tensorflow as tf
def faceToTensor(img):  # Pillow2Tensor
    """
    :param img: PIL格式 图像
    :return: Tensor格式 返回源图像的tensors
    """
    transform = transforms.Compose([
        transforms.Resize([600,600]),
        transforms.ToTensor(),
        transforms.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5])]
    )
    input = transform(img).float().unsqueeze(0)
    return input

def MaskToTensor(img):  # Pillow2Tensor
    """
    :param img: PIL格式 图像
    :return: Tensor格式 返回源图像的tensors
    """
    transform = transforms.Compose([
        transforms.ToTensor(),
        transforms.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5])]
    )
    input = transform(img).float().unsqueeze(0)
    return input

def unnormalize(tensor, mean, std, inplace=False) :
    """Unnormalize a tensor image with mean and standard deviation.

    Args:
        tensor (Tensor): Tensor image of size (C, H, W) or (B, C, H, W) to be normalized.
        mean (sequence): Sequence of means for each channel.
        std (sequence): Sequence of standard deviations for each channel.
        inplace(bool,optional): Bool to make this operation inplace.

    Returns:
        Tensor: Normalized Tensor image.
    """
    if not isinstance(tensor, torch.Tensor):
        raise TypeError('Input tensor should be a torch tensor. Got {}.'.format(type(tensor)))

    if tensor.ndim < 3:
        raise ValueError('Expected tensor to be a tensor image of size (..., C, H, W). Got tensor.size() = '
                         '{}.'.format(tensor.size()))

    if not inplace:
        tensor = tensor.clone()

    dtype = tensor.dtype
    mean = torch.as_tensor(mean, dtype=dtype, device=tensor.device)
    std = torch.as_tensor(std, dtype=dtype, device=tensor.device)
    if (std == 0).any():
        raise ValueError('std evaluated to zero after conversion to {}, leading to division by zero.'.format(dtype))
    if mean.ndim == 1:
        mean = mean.view(-1, 1, 1)
    if std.ndim == 1:
        std = std.view(-1, 1, 1)
    tensor.mul_(std).add_(mean)
    return tensor


def tensor2pillow(input_tensor):
    """
    将tensor保存为pillow
    :param input_tensor: 要保存的tensor
    """
    assert (len(input_tensor.shape) == 4 and input_tensor.shape[0] == 1)
    # 复制一份
    input_tensor = input_tensor.clone().detach()
    # 到cpu
    input_tensor = input_tensor.to(torch.device('cpu'))
    # 反归一化
    input_tensor = unnormalize(input_tensor,[0.5,0.5,0.5],[0.5,0.5,0.5])
    # 去掉批次维度
    input_tensor = input_tensor.squeeze()
    # 从[0,1]转化为[0,255]，再从CHW转为HWC，最后转为numpy
    input_tensor = input_tensor.mul_(255).add_(0.5).clamp_(0, 255).permute(1, 2, 0).type(torch.uint8).numpy()
    # 转成pillow
    im = Image.fromarray(input_tensor)
    return im




def tf_integral(x,a):
	return 0.5*(x*tf.sqrt(x**2+a)+a*tf.log(tf.abs(x+tf.sqrt(x**2+a))))
def tf_pre_parabol(x,par):
	x = x-450.
	prev = 2.*par*(tf_integral(tf.abs(x),0.25/(par**2))-tf_integral(0,0.25/(par**2)))
	return prev+450.
def projector(param, ph, logo):
    '''Apply off-plane transformations to the sticker images
    param: parabola rate of the off-plane parabolic tranformation, rank 2 tensor with shape [N, 1]
    ph:angle of the off-plane rotation, rank 2 tensor with shape [N, 1]
    logo: rank 4 tensor with format NHWC and shape [N, 400, 900, 3]

    return: rank 4 tensor with format NHWC and shape [N, 900, 900, 3]
    '''
    right_cumsum = tf.transpose(
        tf.pad(tf.cumsum(logo[:, :, 450:], axis=2), tf.constant([[0, 0], [0, 0], [1, 0], [0, 0]])), [0, 2, 1, 3])
    left_cumsum = tf.transpose(
        tf.pad(tf.cumsum(logo[:, :, :450][:, :, ::-1], axis=2), tf.constant([[0, 0], [0, 0], [1, 0], [0, 0]])),
        [0, 2, 1, 3])

    anchors = tf.expand_dims(tf.cast(tf.round(tf.clip_by_value( \
        tf_pre_parabol(tf.expand_dims(tf.constant(np.arange(450, 901, dtype=np.float32)), 0), \
                       param) - 450., 0, 450.)), tf.int32), 2)
    anch_inds = tf.tile(tf.expand_dims(tf.expand_dims(tf.range(tf.shape(param)[0]), 1), 2), [1, 451, 1])
    new_anchors = tf.concat([anch_inds, anchors], 2)

    anchors_div = tf.expand_dims(tf.cast(tf.clip_by_value(anchors[:, 1:] - anchors[:, :-1], 1, 900), tf.float32), 3)
    right_anchors_cumsum = tf.gather_nd(right_cumsum, new_anchors)
    right_anchors_diffs = right_anchors_cumsum[:, 1:] - right_anchors_cumsum[:, :-1]
    right = right_anchors_diffs / anchors_div
    left_anchors_cumsum = tf.gather_nd(left_cumsum, new_anchors)
    left_anchors_diffs = left_anchors_cumsum[:, 1:] - left_anchors_cumsum[:, :-1]
    left = left_anchors_diffs / anchors_div

    tmp_result = tf.transpose(tf.concat([left[:, ::-1], right], axis=1), [0, 2, 1, 3])

    cumsum = tf.pad(tf.cumsum(tmp_result, axis=1), tf.constant([[0, 0], [1, 0], [0, 0], [0, 0]]))

    angle = tf.expand_dims(np.pi / 180. * ph, 2)

    z = param * tf.constant((np.arange(900, dtype=np.float32) - 449.5) ** 2)
    z_tile = tf.tile(tf.expand_dims(z, 1), tf.constant([1, 901, 1]))

    y_coord = tf.constant(np.arange(-250, 651, dtype=np.float32))
    y_tile = tf.tile(tf.expand_dims(tf.expand_dims(y_coord, 1), 0), [tf.shape(param)[0], 1, 900])

    y_prev = (y_tile + z_tile * tf.sin(-angle)) / tf.cos(angle)
    y_round = tf.cast(tf.round(tf.clip_by_value(y_prev, 0, 400.)), tf.int32)
    y_div = tf.clip_by_value(y_round[:, 1:] - y_round[:, :-1], 1, 900)

    x_coord = tf.constant(np.arange(900, dtype=np.int32))
    x_tile = tf.tile(tf.expand_dims(tf.expand_dims(x_coord, 0), 0), [tf.shape(param)[0], 901, 1])

    b_coord = tf.tile(tf.expand_dims(tf.expand_dims(tf.range(tf.shape(param)[0]), 1), 2), [1, 901, 900])

    indices = tf.stack([b_coord, y_round, x_tile], axis=3)

    chosen_cumsum = tf.gather_nd(cumsum, indices)
    chosen_cumsum_diffs = chosen_cumsum[:, 1:] - chosen_cumsum[:, :-1]
    final_results = tf.clip_by_value(chosen_cumsum_diffs / tf.expand_dims(tf.cast(y_div, tf.float32), 3), 0., 1.)

    return final_results